import sys, os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
import json
import csv
import argparse
from preprossData.utils.crop_padding import crop_resize_save
from tqdm import tqdm
def paser_label(diagnosis_code):
    """
    将诊断代码解析为文本描述：
    0,1,2,3,4 分别对应：
      - normal
      - mild Diabetic Retinopathy
      - moderate Diabetic Retinopathy
      - severe Diabetic Retinopathy
      - proliferative Diabetic Retinopathy
    """
    mapping = {
        '0': "normal",
        '1': "mild Diabetic Retinopathy",
        '2': "moderate Diabetic Retinopathy",
        '3': "severe Diabetic Retinopathy",
        '4': "proliferative Diabetic Retinopathy"
    }
    result = mapping.get(str(diagnosis_code))
    if result is None:
        raise ValueError(f"Unknown diagnosis code: {diagnosis_code}")
    return result

def anylis_csv(csv_path):
    """
    加载 CSV 文件，文件中包含两列：id_code, diagnosis，
    返回一个字典，键为 id_code，值为 diagnosis（诊断代码）。
    """
    result = {}
    with open(csv_path, newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            id_code = row['id_code'].strip()
            diagnosis = row['diagnosis'].strip()
            result[id_code] = diagnosis
    return result

def gather_data(data_path, tar_path, prefix='aptos'):
    """
    对数据集进行预处理：
      1. 建立目标目录 tar_path（包含 images/ 与 split/ 文件夹）。
      2. 对 train 与 test 分别处理：
         - 从 {data_path}/{split}_images 读取所有图像，
         - 若 split 为 train，则加载对应的 {split}.csv 获取 id_code 与 diagnosis 映射，
           否则 test 中诊断文本置为空字符串。
         - 对每张图像进行裁剪+resize（调用 crop_resize_save），
         - 新图像命名为 {prefix}_{原图 id}.png，
         - 保存预处理后的图像至 tar_path/images 下，
         - 将标注信息写入 data_dict，同时记录划分信息到 split_dict。
      3. 将 split_dict 保存到 tar_path/split/official.json，
         将 data_dict 保存到 tar_path/annotations.json。
    
    参数：
        data_path (str): 原始数据集根目录，要求包含 train_images, test_images, train.csv, test.csv
        tar_path (str): 预处理后数据存放目录
        prefix (str): 新图像文件名前缀，默认为 "aptos"
    
    返回：
        dict: 生成的标注信息字典
    """
    os.makedirs(tar_path, exist_ok=True)
    images_dir = os.path.join(tar_path, 'images')
    os.makedirs(images_dir, exist_ok=True)
    split_dir = os.path.join(tar_path, 'split')
    os.makedirs(split_dir, exist_ok=True)
    
    data_dict = {}
    split_dict = {}
    
    for split in ['train', 'test']:
        src_image_dir = os.path.join(data_path, f"{split}_images")
        image_list = sorted(os.listdir(src_image_dir))
        # 仅对 train 加载 CSV 文件，test 的诊断文本置为空字符串
        if split == "train":
            csv_path = os.path.join(data_path, f"{split}.csv")
            csv_dict = anylis_csv(csv_path)
        else:
            csv_dict = {}
        split_dict[split] = []
        
        for image_name in tqdm(image_list,desc=f"{split} image process",unit='images'):
            # 假设文件名形如 "12345.jpg"，则 id 为 "12345"
            image_id = os.path.splitext(image_name)[0]
            if split == "train":
                diagnose_code = csv_dict.get(image_id)
                if diagnose_code is None:
                    raise ValueError(f"Image id '{image_id}' not found in CSV file: {csv_path}")
                diagnosis_text = paser_label(diagnose_code)
            else:
                diagnosis_text = ""
            
            new_image_name = f"{prefix}_{image_id}.png"
            tar_image_rel_path = os.path.join('images', new_image_name)
            src_image_path = os.path.join(src_image_dir, image_name)
            dest_image_path = os.path.join(tar_path, tar_image_rel_path)
            
            # 对图像进行裁剪与 resize，并保存为 PNG 格式
            crop_info=crop_resize_save(
                image_path=src_image_path,
                save_path=dest_image_path,
                resize=(224, 224),
                crop_threshold=25
            )
            # 构建标注信息。如果 diagnosis_text 非空，则添加 diagnosis 字段
            annotation_entry = {
                'image_name': new_image_name,
                'image_path': tar_image_rel_path,
                'original_path': os.path.relpath(src_image_path,data_path),
                'crop_info':crop_info
            }
            if diagnosis_text:
                annotation_entry['diagnosis'] = {
                    'text': diagnosis_text
                }
            
            data_dict[new_image_name] = annotation_entry
            split_dict[split].append(new_image_name)
    
    # 将 split 信息保存到 tar_path/split/official.json
    official_split_path = os.path.join(split_dir, 'official.json')
    with open(official_split_path, 'w', encoding='utf-8') as f:
        json.dump(split_dict, f, indent=4)
    
    # 将标注信息保存到 tar_path/annotations.json
    annotations_path = os.path.join(tar_path, 'annotations.json')
    with open(annotations_path, 'w', encoding='utf-8') as f:
        json.dump(data_dict, f, indent=4)
    
    return data_dict

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="APTOS 数据集预处理")
    parser.add_argument("--data_path", type=str, default="../Dataset/public_dataset",
                        help="原始数据集根目录，要求包含 train_images, test_images, train.csv, test.csv")
    parser.add_argument("--tar_path", type=str, default="../Dataset/processsed224_public",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="aptos",
                        help="处理后图片名称的前缀，默认为 'aptos'")
    args = parser.parse_args()

    annotations = gather_data(args.data_path, args.tar_path, prefix=args.prefix)
    print("Preprocessing completed.")
